package middlewares

import (
	"gozeroSingle/internal/svc"
	"gozeroSingle/utils"
	"net/http"

	"github.com/zeromicro/go-zero/core/limit"
	"github.com/zeromicro/go-zero/core/stores/redis"
	"github.com/zeromicro/go-zero/rest"
	"github.com/zeromicro/go-zero/rest/httpx"
)

/*
 +----------------------------------------------------------------------
 + Title        : PeriodLimiterMiddleware、TokenLimiterMiddleware
 + Author       : huwl
 + Version      : 1.0.0
 + Initial-Time : 2025/11/04 22:36:33
 + Last-time    : 这个文件最后修改的时间 + 修改人的名称
 + Desc         : 单级限流器
 +----------------------------------------------------------------------
*/

/**
 * 滑动窗口限流器
 * 	1.在固定时间窗口内统计请求次数，超限则拒绝
 * 	2.需要强行限制瞬时请求量，例如防止数据库被刷
 * 	3.固定窗口，规则严格
 * 	4.如果你的场景需要简单粗暴地在单位时间内封顶请求数，比如某个关键接口的严格QPS上限，滑动窗口 (periodlimit) 很合适
 *
 * @param http.HandlerFunc next ------------------------------必传
 * @return http.HandlerFunc
 * @author huwl
 */
func PeriodLimiterMiddleware(svcCtx *svc.ServiceContext) rest.Middleware {
	return func(next http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			// 组装redis配置
			rc := redis.RedisConf{}
			switch svcCtx.Config.Env {
			case "dev":
				rc.Host = svcCtx.Config.RedisConf.Dev.Single.Host
				rc.Pass = svcCtx.Config.RedisConf.Dev.Single.Password
			case "test":
				rc.Host = svcCtx.Config.RedisConf.Test.Single.Host
				rc.Pass = svcCtx.Config.RedisConf.Test.Single.Password
			case "prod":
				rc.Host = svcCtx.Config.RedisConf.Prod.Single.Host
				rc.Pass = svcCtx.Config.RedisConf.Prod.Single.Password
			}

			if svcCtx.Config.RedisConf.Mode == "cluster" {
				rc.Type = "cluster"
			} else {
				rc.Type = "node"
			}

			// 创建全局限流器
			limiter := limit.NewPeriodLimit(
				svcCtx.Config.LimiterConf.GlobalLimiter.Period,
				svcCtx.Config.LimiterConf.GlobalLimiter.Quota,
				redis.MustNewRedis(rc),
				"global_limit")
			key := r.URL.Path
			code, err := limiter.Take(key)
			if err != nil {
				httpx.OkJsonCtx(r.Context(), w, utils.ErrorOutput(err.Error()))
				return
			}
			switch code {
			case limit.Allowed:
				next(w, r)
			case limit.HitQuota, limit.OverQuota:
				httpx.OkJsonCtx(r.Context(), w, utils.ErrorOutput("人气太火爆了，请稍等", []int{http.StatusTooManyRequests}...))
				return
			}
		}
	}
}

/**
 * token限流器
 * 	1.以恒定速率向桶中添加令牌，请求获取到令牌才可通过
 * 	2.限制平均速率，同时允许一定程度的突发流量
 * 	3.允许流量突发，更具弹性
 * 	4.如果希望系统在空闲时能"攒下"一些容量，以应对短暂的流量峰值，比如网关、API入口等场景，令牌桶 (tokenlimit) 是更好的选择
 *
 * @param http.HandlerFunc next ------------------------------必传
 * @return http.HandlerFunc
 * @author huwl
 */
func TokenLimiterMiddleware(svcCtx *svc.ServiceContext) rest.Middleware {
	return func(next http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			// 组装redis配置
			rc := redis.RedisConf{}
			switch svcCtx.Config.Env {
			case "dev":
				rc.Host = svcCtx.Config.RedisConf.Dev.Single.Host
				rc.Pass = svcCtx.Config.RedisConf.Dev.Single.Password
			case "test":
				rc.Host = svcCtx.Config.RedisConf.Test.Single.Host
				rc.Pass = svcCtx.Config.RedisConf.Test.Single.Password
			case "prod":
				rc.Host = svcCtx.Config.RedisConf.Prod.Single.Host
				rc.Pass = svcCtx.Config.RedisConf.Prod.Single.Password
			}

			if svcCtx.Config.RedisConf.Mode == "cluster" {
				rc.Type = "cluster"
			} else {
				rc.Type = "node"
			}

			// 创建令牌桶限流器
			// 参数：每秒生成的令牌数(rate)，令牌桶最大容量(burst)
			limiter := limit.NewTokenLimiter(
				svcCtx.Config.LimiterConf.GlobalLimiter.Rate,
				svcCtx.Config.LimiterConf.GlobalLimiter.Burst,
				redis.MustNewRedis(rc),
				"global_limit")
			if limiter.Allow() {
				next(w, r)
			} else {
				httpx.OkJsonCtx(r.Context(), w, utils.ErrorOutput("人气太火爆了，请稍等", []int{http.StatusTooManyRequests}...))
				return
			}
		}
	}
}
